import torch
from norm_layer import get_normalize_layer
from pytorchcv.model_provider import get_model as ptcv_get_model

def load_imagenet_model(model_name,device):
    normalize_layer = get_normalize_layer("imagenet")
    model = None
    if model_name == 'vgg16':
        model = ptcv_get_model("vgg16", pretrained=False)
        checkpoint = torch.load("model_weights/imagenet/vgg16-0865-5ca155da.pth", map_location=device)
    elif model_name == 'resnet50':
        model = ptcv_get_model("resnet50", pretrained=False)
        checkpoint = torch.load("model_weights/imagenet/resnet50-0633-b00d1c8e.pth", map_location=device)
    elif model_name == 'resnext26-m':  # **
        model = ptcv_get_model("resnext26_32x4d", pretrained=False)
        checkpoint = torch.load("model_weights/imagenet/resnext26_32x4d-0746-1011ac35.pth", map_location=device)
    elif model_name == 'seresnet18-m':  # **
        model = ptcv_get_model("seresnet18", pretrained=False)
        checkpoint = torch.load("model_weights/imagenet/seresnet18-0961-022123a5.pth", map_location=device)
    elif model_name == 'scnet50-m':   # **
        model = ptcv_get_model("scnet50", pretrained=False)
        checkpoint = torch.load("model_weights/imagenet/scnet50-0547-18741240.pth", map_location=device)

    model.load_state_dict(checkpoint)
    return torch.nn.Sequential(normalize_layer, model)

def load_user_model(device, path):
    normalize_layer = get_normalize_layer("imagenet")
    model = get_model()
    checkpoint = torch.load(path, map_location=device)
    model.load_state_dict(checkpoint)
    return torch.nn.Sequential(normalize_layer, model)
